import os
import os.path as osp
import numpy as np
import cv2
import matplotlib.pyplot as plt
#根据生成的预测box的txt进行可视化

#det_file 为对应类别的txt文件，格式如下
# 图片名称 置信度 四个点的坐标
# 000089 0.9956 302.9024 180.4394 296.8925 155.5844 282.2737 159.1192 288.2836 183.9742
# 000089 0.9954 271.3734 95.2678 269.6444 70.9205 259.2952 71.6554 261.0242 96.0027
def main(det_file,img_dir,save_dir):
    detections_dict = genInfoDict(det_file)
    with open(det_file, 'r') as f:
        lines = f.readlines()
    splitlines = [x.strip().split(' ') for x in lines]
    image_ids = list(set([x[0] for x in splitlines]))
    for i,image_id in enumerate(image_ids):
        print(i," finished")
        image_path = osp.join(img_dir,"{}.jpg".format(image_id))
        img = cv2.imread(image_path)
        quads = detections_dict[image_id]
        drawRect(img,quads,save_dir,image_id)
    
    
def drawRect(img,quads,save_dir,image_id):
    
    for quad in quads:
        score = float(quad[0])
        #修改显示的置信度
        if score >0.3:
            quad[1:] = [int(float(t)) for t in quad[1:]]
            x0n,y0n,x1n,y1n,x2n,y2n,x3n,y3n = quad[1:]
            cv2.line(img,(x0n, y0n),(x1n, y1n), (0, 0, 255), thickness=2, lineType=8)
            cv2.line(img,(x1n, y1n),(x2n, y2n), (255, 0, 0), thickness=2, lineType=8)
            cv2.line(img,(x2n, y2n),(x3n, y3n), (0, 0, 255), thickness=2, lineType=8)
            cv2.line(img,(x0n, y0n), (x3n, y3n), (255,0, 0), thickness=2, lineType=8)
            
    cv2.imwrite(osp.join(save_dir,"{}.jpg".format(image_id)), img)



    
    
    
def genInfoDict(det_file):
    with open(det_file, 'r') as f:
        lines = f.readlines()
    splitlines = [x.strip().split(' ') for x in lines]
    image_ids = tuple(set([x[0] for x in splitlines]))

    lists = [[] for x in range(len(image_ids))]
    #初始化对应图片的box参数的dict
    detections_dict =dict(zip(image_ids,lists))

    for splitline in splitlines:
        image_id = splitline[0]
#         print(image_id )
        info = splitline[1:]

        detections_dict[image_id].append(info)
    return detections_dict

        
    


det_file = "Outputs/vgg16_nwpu/Oct09-08-53-06_6804d4778686_step/test/model_step37499/comp4_0e91a2aa-c344-4af1-82b9-1dc81cc938ad_det_test_airplane.txt"
img_dir = "../_DATASET/NWPU_VHR/AllImages"
save_dir = "Outputs/vgg16_nwpu/Oct09-08-53-06_6804d4778686_step/vis"
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
main(det_file,img_dir,save_dir)